# server.py
import openpyxl
from mcp.server.fastmcp import FastMCP
from typing import Any, Callable, List, Optional, cast
import os
import httpx
import json


# Create an MCP server
mcp = FastMCP("Demo")


# Add an addition tool
@mcp.tool()
def add(a: int, b: int) -> int:
    """Add two numbers"""
    return a + b


# Add a dynamic greeting resource
@mcp.resource("greeting://{name}")
def get_greeting(name: str) -> str:
    """Get a personalized greeting"""
    return f"Hello, {name}!"


@mcp.resource('echo://static')
async def echo_resource():
    # 返回的是，当用户使用这个资源时，资源的内容
    return 'Echo!'


@mcp.tool()
# 第一个参数会被传入上下文对象
async def web_search(query: str) -> str:
    """
    搜索互联网内容

    Args:
        query: 要搜索内容

    Returns:
        搜索结果的总结
    """
    async with httpx.AsyncClient() as client:
        response = await client.post(
            'https://open.bigmodel.cn/api/paas/v4/tools',
            headers={'Authorization': '3f9fa9317d1fc918fcc2e3e8a73afdf7.0dqwi4kw0eaUoUvi'},
            json={
                'tool': 'web-search-pro',
                'messages': [
                    {'role': 'user', 'content': query}
                ],
                'stream': False
            }
        )

        res_data = []
        for choice in response.json()['choices']:
            for message in choice['message']['tool_calls']:
                search_results = message.get('search_result')
                if not search_results:
                    continue
                for result in search_results:
                    res_data.append(result['content'])

        return_data = '\n\n\n'.join(res_data)

        return return_data


@mcp.tool()
async def fetch_weather(city: str) -> str:
    """Fetch current weather for a city"""
    async with httpx.AsyncClient() as client:
        response = await client.get(f"https://api.weather.com/{city}")
        return response.text

#
# @mcp.tool()
# async def query_meta(item_class: str) -> Optional[dict[str, Any]]:
#     '''
#     查询实体类元信息，如实体类名称、属性列表等
#     参数：
#         item_class->str：要查询的实体类名称
#     '''
#     print(f"查询接口: {item_class}")
#     from dwf_client import client
#     url = f"http://172.21.11.79:9090/dwf/v1/meta/entities/{item_class}"
#     r = client.request(url, 'GET')
#     return r

#
# @mcp.tool()
# async def query_data(item_class: str, condition: str) -> Optional[dict[str, Any]]:
#     '''
#     查询实体类数据
#     参数：
#         item_class->str：要查询的实体类名称
#         condition->str: 查询条件
#     '''
#     print(f"查询接口: {item_class}, {condition}")
#     from dwf_client import client
#     url = f"http://172.21.11.79:9090/dwf/v1/omf/entities/{item_class}/objects"
#     r = client.request(url, 'POST', data={"condition": ""})
#     return r




@mcp.prompt()
def review_code(code: str) -> str:
    return f"Please review this code:\n\n{code}"


@mcp.tool()
def search_desktop_file(filename: str) -> list[str]:
    """
    在桌面上搜索与给定文件名部分匹配的文件路径（支持模糊匹配）。

    参数:
    - filename: 要搜索的文件名关键词，支持模糊匹配（不区分大小写）。

    返回:
    - 所有匹配文件的完整路径列表。
    """
    desktop_path = os.path.join(os.environ["USERPROFILE"], "Desktop")
    matched_files = []
    for root, _, files in os.walk(desktop_path):
        for file in files:
            if filename.lower() in file.lower():
                matched_files.append(os.path.join(root, file))
    return matched_files


@mcp.tool()
def read_xlsx_file(filepath: str) -> List[List[str]]:
    """
    读取指定路径的 xlsx 文件内容。

    参数:
    - filepath: 要读取的 xlsx 文件的完整路径。

    返回:
    - 二维列表，表示 Excel 表格中的所有单元格内容。
      每个内层列表代表一行，行内是该行的单元格字符串。

    异常:
    - 如果文件不存在或不是 .xlsx 文件，会抛出异常。
    """
    if not os.path.isfile(filepath):
        raise FileNotFoundError(f"文件不存在: {filepath}")
    if not filepath.lower().endswith(".xlsx"):
        raise ValueError("只支持读取 .xlsx 格式的文件")

    wb = openpyxl.load_workbook(filepath, data_only=True)
    sheet = wb.active  # 默认读取第一个工作表
    data = []

    for row in sheet.iter_rows(values_only=True):
        data.append([str(cell) if cell is not None else "" for cell in row])

    return data

@mcp.tool()
def read_file_content(path: str) -> str:
    """
    读取指定文件的文本内容。

    Args::
    - path: 要读取的文件的完整路径。

    Return:
    - 文件的文本内容字符串。如果文件不存在或读取失败，将抛出异常。
    """
    if not os.path.isabs(path):
        path = os.path.join(os.environ["USERPROFILE"], "Desktop", path)
    if not os.path.exists(path):
        raise FileNotFoundError(f"文件不存在：{path}")
    with open(path, 'r', encoding='utf-8', errors='ignore') as f:
        return f.read()


@mcp.tool()
def append_to_file(path: str, content: str) -> str:
    """
    向指定文件中追加一段文本内容（以换行结尾）。

    参数:
    - path: 要写入的文件路径。
    - content: 要追加写入的文本内容。

    返回:
    - 写入结果字符串。如果成功，返回 "success"。
    """
    if not os.path.isabs(path):
        path = os.path.join(os.environ["USERPROFILE"], "Desktop", path)
    with open(path, 'a', encoding='utf-8') as f:
        f.write(content + '\n')
    return "success"


@mcp.tool()
async def image_generation(image_prompt: str):
    """
    生成图片，绘制图片
    :param image_prompt: 图片描述，需要是英文
    :return: 图片保存到的本地路径
    """
    async with httpx.AsyncClient() as client:
        data = {'data': [image_prompt, 0, True, 512, 512, 3]}

        # 创建生成图片任务
        response1 = await client.post(
            'https://black-forest-labs-flux-1-schnell.hf.space/call/infer',
            json=data,
            headers={"Content-Type": "application/json"}
        )

        # 解析响应获取事件 ID
        response_data = response1.json()
        event_id = response_data.get('event_id')

        if not event_id:
            return '无法获取事件 ID'

        # 通过流式的方式拿到返回数据
        url = f'https://black-forest-labs-flux-1-schnell.hf.space/call/infer/{event_id}'
        full_response = ''
        async with client.stream('GET', url) as response2:
            async for chunk in response2.aiter_text():
                full_response += chunk

        return json.loads(full_response.split('data: ')[-1])[0]['url']

if __name__ == "__main__":

    mcp.run()

